// #define __KERNEL__
// #define MODULE
#include <linux/module.h>
#include <linux/kernel.h>
// #include <linux/cons.h>
#define __NO_VERSION__
#include <rtai.h>
#include <rtai_sched.h>
#include <rtai_shm.h>
#include "rtai_shmem.h"

#define MAX_TASK_NUM 4

// 全局变量
static RT_TASK task_handle;

// 调度器任务函数
static void task_func(long arg)
{
    int i;
    char name[32];
    unsigned long count=0;
    unsigned long value;
    unsigned long cycle_time_ns=1000000;
    unsigned int cycyle_offset_ns;
    RTIME start_time, prev_time;
    int latency, jitter, prev_latency = 0;
    static unsigned int max_latency = 0;
    static unsigned int max_jitter = 0;

    RTIME next_period = rt_get_time() + nano2count(cycle_time_ns);
    prev_time = 0;
    
    while (1) {
        start_time = rt_get_time();
        if (count > 1) {
            latency = count2nano(start_time - prev_time) - cycle_time_ns;
            if (abs(latency) > max_latency) {
                max_latency = abs(latency);
                shmem_write("max_latency", &max_latency);
            }
            
            // 计算jitter (当前latency与上一个latency的差值)
            if(count > 2){
                jitter = abs(latency - prev_latency);
                if (jitter > max_jitter) {
                    max_jitter = jitter;
                    shmem_write("max_jitter", &max_jitter);
                }
            }
            prev_latency = latency;
        }
        prev_time = start_time;

        // 执行所有激活的任务
        for (i = 0; i < MAX_TASK_NUM; i++) {
            sprintf(name, "taskslot%d", i);
            shmem_read(name, &value);
            if(value != 0){
                // 将value转换为函数指针并调用
                void (*func_ptr)(void) = (void (*)(void))value;
                func_ptr();
            }
        }

        shmem_read("cycle_time", &cycle_time_ns);
        shmem_write("cycle_offset", &cycyle_offset_ns);
        rt_sleep_until(next_period);
        next_period += nano2count(cycle_time_ns + cycyle_offset_ns);
        count++;
        shmem_write("count", &count);
    }
}

// 初始化调度器
int init_module(void)
{
    
    int i;
    char name[32];
    unsigned long value = 0;
    // RTIME now = rt_get_time();
    
    // 分配共享内存
    rtai_shmem_init();
    
    value = 1000000;
    shmem_register("cycle_time", TYPE_UDINT, &value);
    value = 0;
    shmem_register("cycle_offset", TYPE_UINT, &value);
    shmem_register("count", TYPE_UDINT, &value);
    shmem_register("max_jitter", TYPE_UINT, &value);
    shmem_register("max_latency", TYPE_UINT, &value);
    for (i = 0; i < MAX_TASK_NUM; i++) {
        sprintf(name, "taskslot%d", i);
        shmem_register(name, TYPE_UDINT, &value);
    }
    
    // 初始化实时任务并绑定到1号CPU核心
    rt_task_init_cpuid(&task_handle, task_func, 0, 4096, 1, 0, NULL, 1);
    rt_task_resume(&task_handle);
    
    // 设置任务为周期性运行 (1ms周期)
    // rt_task_make_periodic(&task_handle, now + 2*1000000, 1000000);

    printk("RTAI Threadripper initialize\n");
    
    return 0;
}

// 清理调度器
void cleanup_module(void)
{
    rt_task_delete(&task_handle);
    rtai_kfree(SHM_KEY);
    printk("RTAI Threadripper exiting\n");
}

MODULE_LICENSE("GPL");
